/*****************************************************
 Copyright Buu Information Security Major
 File Name: secretsender.c
 Author: 佟美怡
 ID: 2023240381036
 Version: 3.0
 Date: 2025-6-17
 Description:
    密码学课程设计完整实现，包含：
    - 基本通信框架
    - 明文传输功能
    - Diffie-Hellman密钥交换
    - 移位加密算法
 *****************************************************/

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>
#include <winsock2.h>
#include <ws2tcpip.h>

#pragma comment(lib, "ws2_32.lib")

/****************************************************
 Macro Definitions:
    PORT: 通信端口号
    MAX_DATA_LEN: 最大数据长度
    BACKLOG: 最大等待连接数
    CODE_SPACE: 编码空间大小
    DH_P: Diffie-Hellman参数p
    DH_A: Diffie-Hellman参数a
 ****************************************************/
#define PORT 27013
#define MAX_DATA_LEN 1024
#define BACKLOG 5
#define CODE_SPACE 127
#define DH_P 23
#define DH_A 5

/****************************************************
 Enum: MsgType
 Description: 协议消息类型定义
 Values:
    SENDPAR: 发送公共参数
    PAROK: 确认公共参数
    MyNum: 发送自己的公开值
    Bye: 结束会话
 ****************************************************/
typedef enum {
    SENDPAR = 1,
    PAROK,
    MyNum,
    Bye
} MsgType;

/****************************************************
 Struct: DH_Protocol
 Description: 协议数据结构定义
 Members:
    version: 协议版本
    type: 消息类型
    p: 参数p
    a: 参数a
    Y: 共享给对方的数值
 ****************************************************/
typedef struct {
    unsigned char version;
    MsgType type;
    unsigned short p;
    unsigned short a;
    unsigned short Y;
} DH_Protocol;

/****************************************************
 Function: print_usage
 Parameters: None
 Return: None
 Description: 打印程序使用说明
 ****************************************************/
void print_usage() {
    printf("Parameter error.\n");
}

/****************************************************
 Function: mod_exp
 Parameters:
    base: 基数
    exp: 指数
    mod: 模数
 Return: 计算结果 (base^exp) % mod
 Description: 计算模幂运算
 ****************************************************/
unsigned short mod_exp(unsigned short base, unsigned short exp, unsigned short mod) {
    unsigned short result = 1;
    base = base % mod;
    
    while (exp > 0) {
        if (exp % 2 == 1)
            result = (result * base) % mod;
        exp = exp >> 1;
        base = (base * base) % mod;
    }
    return result;
}

/****************************************************
 Function: dh_key_exchange
 Parameters:
    sock: 通信socket
    is_server: 是否为服务器端(1:服务器, 0:客户端)
 Return: 协商出的密钥
 Description: 执行Diffie-Hellman密钥交换协议
 ****************************************************/
unsigned short dh_key_exchange(SOCKET sock, int is_server) {
    DH_Protocol protocol;
    unsigned short private_key, public_key, shared_key;
    
    // 生成私钥(1 < private_key < p-1)
    srand((unsigned int)time(NULL));
    private_key = (rand() % (DH_P - 2)) + 1;
    
    // 设置协议版本
    protocol.version = 1;
    
    if (is_server) {
        // 服务器端流程
        // 接收客户端公共参数
        recv(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 发送确认
        protocol.type = PAROK;
        send(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 计算并发送自己的公开值
        public_key = mod_exp(DH_A, private_key, DH_P);
        protocol.type = MyNum;
        protocol.Y = public_key;
        send(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 接收客户端的公开值
        recv(sock, (char*)&protocol, sizeof(protocol), 0);
        shared_key = mod_exp(protocol.Y, private_key, DH_P);
    } else {
        // 客户端流程
        // 发送公共参数
        protocol.type = SENDPAR;
        protocol.p = DH_P;
        protocol.a = DH_A;
        send(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 接收确认
        recv(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 计算并发送自己的公开值
        public_key = mod_exp(DH_A, private_key, DH_P);
        protocol.type = MyNum;
        protocol.Y = public_key;
        send(sock, (char*)&protocol, sizeof(protocol), 0);
        
        // 接收服务器的公开值
        recv(sock, (char*)&protocol, sizeof(protocol), 0);
        shared_key = mod_exp(protocol.Y, private_key, DH_P);
    }
    
    return shared_key % DH_P; // 最终密钥
}

/****************************************************
 Function: encrypt_data
 Parameters:
    data: 待加密数据
    key: 加密密钥
    output: 加密输出缓冲区(十六进制格式)
 Return: None
 Description: 使用移位加密算法加密数据
 ****************************************************/
void encrypt_data(const char* data, int key, char* output) {
    int i;
    char temp[16];
    
    strcpy(output, "Hex:");
    for (i = 0; i < strlen(data); i++) {
        unsigned char c = (data[i] + key) % CODE_SPACE;
        sprintf(temp, " %02x", c);
        strcat(output, temp);
    }
}

/****************************************************
 Function: decrypt_data
 Parameters:
    hex_data: 十六进制格式的加密数据
    key: 解密密钥
    output: 解密输出缓冲区
 Return: None
 Description: 使用移位加密算法解密数据
 ****************************************************/
void decrypt_data(const char* hex_data, int key, char* output) {
    int i, val;
    char hex[3] = {0};
    const char* ptr = hex_data + 4; // 跳过"Hex:"前缀
    
    i = 0;
    while (sscanf(ptr, "%2x", &val) == 1) {
        output[i++] = (val - key + CODE_SPACE) % CODE_SPACE;
        ptr += 3; // 移动指针跳过空格和两个十六进制字符
    }
    output[i] = '\0';
}

/****************************************************
 Function: init_server_socket
 Parameters:
    server_socket: 服务器socket指针
 Return: 成功返回0，失败返回-1
 Description: 初始化服务器socket，绑定端口并开始监听
 ****************************************************/
int init_server_socket(SOCKET* server_socket) {
    struct sockaddr_in server_addr;
    
    // 创建socket
    *server_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (*server_socket == INVALID_SOCKET) {
        printf("socket failed with error: %d\n", WSAGetLastError());
        return -1;
    }

    // 设置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(PORT);

    // 绑定socket
    if (bind(*server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        printf("bind failed with error: %d\n", WSAGetLastError());
        closesocket(*server_socket);
        return -1;
    }

    // 开始监听
    if (listen(*server_socket, BACKLOG) == SOCKET_ERROR) {
        printf("listen failed with error: %d\n", WSAGetLastError());
        closesocket(*server_socket);
        return -1;
    }

    return 0;
}

/****************************************************
 Function: run_server
 Parameters: None
 Return: 程序执行状态(0:成功, 1:失败)
 Description: 服务器模式主函数
 ****************************************************/
int run_server() {
    SOCKET server_socket, client_socket;
    struct sockaddr_in client_addr;
    int client_addr_len = sizeof(client_addr);
    char recv_buf[MAX_DATA_LEN];
    int recv_len;
    unsigned short key;

    printf("waiting connected...");

    // 初始化服务器socket
    if (init_server_socket(&server_socket) != 0) {
        return 1;
    }

    // 接受客户端连接
    client_socket = accept(server_socket, (struct sockaddr*)&client_addr, &client_addr_len);
    if (client_socket == INVALID_SOCKET) {
        printf("accept failed with error: %d\n", WSAGetLastError());
        closesocket(server_socket);
        return 1;
    }
    printf("ok\n");

    // 执行Diffie-Hellman密钥交换
    key = dh_key_exchange(client_socket, 1);
    printf("key=%d\n", key);

    // 接收数据
    recv_len = recv(client_socket, recv_buf, MAX_DATA_LEN, 0);
    if (recv_len > 0) {
        recv_buf[recv_len] = '\0';  // 确保字符串结束
        
        // 判断是否为密文(Hex:开头)
        if (strncmp(recv_buf, "Hex:", 4) == 0) {
            char plaintext[MAX_DATA_LEN];
            decrypt_data(recv_buf, key, plaintext);
            printf("received cipher(hex) <%s>\n", recv_buf + 4);
            printf("plain text <%s>\n", plaintext);
        } else {
            printf("received data <%s>\n", recv_buf);
        }
    }

    // 关闭socket
    closesocket(client_socket);
    closesocket(server_socket);
    
    return 0;
}

/****************************************************
 Function: connect_to_server
 Parameters:
    ip: 服务器IP地址
    client_socket: 客户端socket指针
 Return: 成功返回0，失败返回-1
 Description: 连接指定IP的服务器
 ****************************************************/
int connect_to_server(const char* ip, SOCKET* client_socket) {
    struct sockaddr_in server_addr;

    // 创建socket
    *client_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (*client_socket == INVALID_SOCKET) {
        printf("socket failed with error: %d\n", WSAGetLastError());
        return -1;
    }

    // 设置服务器地址
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(PORT);
    if (inet_pton(AF_INET, ip, &server_addr.sin_addr) <= 0) {
        printf("Invalid address/Address not supported\n");
        return -1;
    }

    // 连接服务器
    if (connect(*client_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        printf("connect failed with error: %d\n", WSAGetLastError());
        closesocket(*client_socket);
        return -1;
    }

    return 0;
}

/****************************************************
 Function: run_client
 Parameters:
    ip: 服务器IP地址
    option: 选项参数(-n或-m)
    data: 要发送的数据
 Return: 程序执行状态(0:成功, 1:失败)
 Description: 客户端模式主函数
 ****************************************************/
int run_client(const char* ip, const char* option, const char* data) {
    SOCKET client_socket;
    int send_len;
    unsigned short key;
    char encrypted[MAX_DATA_LEN * 3]; // 十六进制格式需要更多空间

    // 检查数据长度
    if (strlen(data) == 0) {
        printf("Error:Data length is zero.\n");
        return 1;
    }

    // 检查加密选项
    if (strcmp(option, "-n") != 0 && strcmp(option, "-m") != 0) {
        printf("Parameter error.\n");
        return 1;
    }

    // 连接到服务器
    if (connect_to_server(ip, &client_socket) != 0) {
        return 1;
    }
    printf("connect %s: ok\n", ip);

    if (strcmp(option, "-m") == 0) {
        // 密文模式
        key = dh_key_exchange(client_socket, 0);
        printf("key=%d\n", key);
        
        // 加密数据
        encrypt_data(data, key, encrypted);
        printf("Data length:%d\n", (int)strlen(data));
        printf("Plain <%s>\n", data);
        printf("Cypher(hex)%s\n", encrypted + 3); // 去掉"Hex"
        
        // 发送加密数据
        printf("send data %s : ", encrypted);
        send_len = send(client_socket, encrypted, strlen(encrypted), 0);
    } else {
        // 明文模式
        printf("Data length:%d\n", (int)strlen(data));
        printf("send data <%s> : ", data);
        send_len = send(client_socket, data, (int)strlen(data), 0);
    }

    if (send_len == SOCKET_ERROR) {
        printf("failed with error: %d\n", WSAGetLastError());
    } else {
        printf("ok\n");
    }

    // 关闭socket
    closesocket(client_socket);
    
    return 0;
}

/****************************************************
 Function: main
 Parameters:
    argc: 参数个数
    argv: 参数数组
 Return: 程序执行状态(0:成功, 1:失败)
 Description: 程序主入口函数
 ****************************************************/
int main(int argc, char* argv[]) {
    WSADATA wsaData;
    
    // 初始化Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        printf("WSAStartup failed.\n");
        return 1;
    }

    // 参数检查
    if (argc < 2) {
        print_usage();
        WSACleanup();
        return 1;
    }

    // 解析参数
    if (strcmp(argv[1], "-s") == 0) {
        // 服务器模式
        if (argc != 2) {
            printf("Parameter error.\n");
            WSACleanup();
            return 1;
        }
        run_server();
    } 
    else if (strcmp(argv[1], "-c") == 0) {
        // 客户端模式
        if (argc != 5) {
            printf("Parameter error.\n");
            WSACleanup();
            return 1;
        }
        run_client(argv[2], argv[3], argv[4]);
    } 
    else {
        printf("Parameter error.\n");
        WSACleanup();
        return 1;
    }

    // 清理Winsock
    WSACleanup();
    return 0;
}